from dataclasses import dataclass
from typing import Iterable, List, Optional

from collections import defaultdict

import numpy as np
from torch.utils.data import Sampler as TorchSampler


@dataclass
class SamplerConfig:
    """
    Sampler configuration.

    Attributes:
        enable (bool): Whether the sampler is enabled.
        nb_utterances (Optional[int]): Number of utterances.
        nb_speakers (Optional[int]): Number of speakers.
        nb_samples_per_spk (Optional[int]): Number of samples per speaker.
        create_contrastive_pairs (bool): Whether to create contrastive pairs.
        prevent_class_collisions (bool): Whether to prevent class collisions.
        randomize_at_each_epoch (bool): Whether to randomize the starting indices at each epoch.
    """

    enable: bool = True
    nb_utterances: Optional[int] = None
    nb_speakers: Optional[int] = None
    nb_samples_per_spk: Optional[int] = None
    create_contrastive_pairs: bool = False
    prevent_class_collisions: bool = False
    randomize_at_each_epoch: bool = False


class Sampler(TorchSampler):
    """
    PyTorch sampler to generate specific indices for data sampling.

    Attributes:
        labels (List[int]): Dataset labels.
        batch_size (int): Batch size.
        config (SamplerConfig): Sampler configuration.
        seed (int): Seed for reproducibility.
        epoch (int): Current epoch.
        count (int): Total number of indices.
    """

    def __init__(
        self,
        labels: List[int],
        batch_size: int,
        config: SamplerConfig,
        seed: int = 0,
    ):
        """
        Initialize a Sampler object.

        Args:
            labels (List[int]): Dataset labels.
            batch_size (int): Batch size.
            config (SamplerConfig): Sampler configuration.
            seed (int): Seed for reproducibility. Defaults to 0.
        """
        self.labels = labels
        self.batch_size = batch_size
        self.config = config
        self.seed = seed

        self.epoch = 0
        self.count = 0

    def __len__(self) -> int:
        """
        Get the number of indices generated by the sampler.

        Returns:
            int: Length of the sampler.
        """
        return self.count

    def set_epoch(self, epoch: int):
        """
        Set the current epoch.

        Args:
            epoch (int): Current epoch.

        Returns:
            None
        """
        self.epoch = epoch

    def __iter__(self) -> Iterable[int]:
        """
        Generate indices.

        Returns:
            Iterable[int]: Iterable of indices.
        """
        if self.config.randomize_at_each_epoch:
            self.seed = self.seed + self.epoch
        rng = np.random.default_rng(seed=self.seed)

        total_nb_samples = len(self.labels)
        
        nb_samples = total_nb_samples
        if self.config.nb_utterances:
            nb_samples = min(self.config.nb_utterances, total_nb_samples)

        indices = rng.choice(total_nb_samples, size=nb_samples, replace=False)

        # Create list of utterances for each speaker
        spk_to_utterances = defaultdict(list)
        for i in indices:
            spk_to_utterances[self.labels[i]].append(int(i))
        speakers = list(spk_to_utterances.keys())
        speakers.sort()

        # Create indexes by providing pairs if create_contrastive_pairs
        # is set to True
        x = []
        y = []
        for i, speaker in enumerate(speakers):
            if self.config.nb_speakers and i >= self.config.nb_speakers:
                break

            utterances = spk_to_utterances[speaker]

            nb_utt = len(utterances)

            if self.config.nb_samples_per_spk:
                nb_utt = min(nb_utt, self.config.nb_samples_per_spk)

            if self.config.create_contrastive_pairs:
                nb_utt = nb_utt - nb_utt % 2
                for idx in range(0, nb_utt, 2):
                    x.append((utterances[idx], utterances[idx + 1]))
                    y.append(i)
                    x.append((utterances[idx + 1], utterances[idx]))
                    y.append(i)
            else:
                for idx in range(nb_utt):
                    x.append(utterances[idx])
                    y.append(i)

        # Shuffle indices and avoid having two samples of the same speaker
        # in the same batch if prevent_class_collisions is set to True
        x_ = []
        y_ = []
        for i in rng.permutation(len(y)):
            batch_start_i = len(y_) - len(y_) % self.batch_size
            if self.config.prevent_class_collisions and y[i] in y_[batch_start_i:]:
                continue

            x_.append(x[i])
            y_.append(y[i])

        self.count = len(x_) - len(x_) % self.batch_size
        return iter(x_[: self.count])
